'''
Created on Oct 7, 2009

@author: yanghoch
'''
import operator

class Isoform( object ):
    def __init__( self, Name ):
        self.Name = Name
        self.strand = '+'
        self.Seq = ""
        self.exons = []
        self.junctions = []
        self.exp = 1
    def get_exons( self, line, sep1, sep2 ):
        ''' 
        sep1 is to separate each exons 
        sep2 is to separate each start and end of exons
        '''
        exons_strs = line.split( sep1 )
        for e in exons_strs:
            e = e.strip()
            if sep2 in e:
                exon = e.split( sep2 )
                exon = ( int( exon[0] ), int( exon[1] ) )
                self.exons.append( exon )
        if ( len( self.exons ) > 1 ):
            self.exons.sort\
                ( cmp = None, key = operator.itemgetter( 0 ), reverse = False )
            self.__get_junctions__()

        # get junctions
    def __get_junctions__( self ):
        ''' Translate self.exons into junctions info'''
        self.junctions = []
        for j in range( len( self.exons ) - 1 ):
            junc = ( self.exons[j][1], self.exons[j + 1][0] )
            self.junctions.append( junc )

class Isoforms( object ):
    INameSep = ':'
    def __init__( self, Name ):
        self.Isoforms = dict()
    def get_exons_info( self, exonsFN, sep1, sep2 ):
        ''' read exons information from the given file'''
        exonsF = open( exonsFN )
        for line in exonsF.readlines():
            # change this for different format
            IsoformN_exons = line.split( self.INameSep )
            if( len( IsoformN_exons ) == 2 ):
                [IsoformN , exons] = IsoformN_exons
                if not ( IsoformN in self.Isoforms ):
                    I = Isoform( IsoformN )
                    if IsoformN[1] == 'R':
                        I.strand = '-'
                    self.Isoforms[IsoformN] = I
                self.Isoforms[IsoformN].get_exons( exons, sep1, sep2 )
            else:
                print IsoformN_exons # unexpected format
        exonsF.close()

    def get_expression_info( self, expFN, sep ):
        ''' function to get TranSEQ output format with
        (1) Isoform Name (2) Expression Level (3) Log LR
        '''
        def roundAtof( floatStr ):
            if 'e' in floatStr:
                return( 0.0 )
            elif'Inf' in floatStr:
                return ( floatStr )
            else:
                return ( float )( floatStr )
        expF = open( expFN )
        expF.readline() # skip the first header line
        for line in expF.readlines():
            # change this for different format
            [IsoformN, expression, logLR] = line.split( sep )
            if not ( IsoformN in self.Isoforms ):
                I = Isoform( IsoformN )
                if IsoformN[1] == 'R':
                    I.strand = '-'
                elif IsoformN[1] == 'F':
                    I.strand = '-'
                else:
                    print "unexpected name", IsoformN
                self.Isoforms[IsoformN] = I
            self.Isoforms[IsoformN].exp = roundAtof( expression )
            self.Isoforms[IsoformN].logLR = roundAtof( logLR )
        expF.close()

    def get_junc_info( self ):
        def add_junctions( JDict, I ):
            for j in I.junctions:
                juncN = str( j[0] ) + '_' + str( j[1] )
                if not ( juncN in JDict ):
                    JDict[juncN] = I.exp
                else:
                    JDict[juncN] += I.exp
        self.F_Juncs = dict()
        self.R_Juncs = dict()
        for key in self.Isoforms:
            if self.Isoforms[key].strand == '+':
                add_junctions( self.F_Juncs, self.Isoforms[key] )
            else:
                add_junctions( self.R_Juncs, self.Isoforms[key] )

    def summary_Juncs_info( self ):
        ''' sort junctions and report their exp frequencies'''
        def sort_junc_dict( JuncDic ):
            sortedJuncs = []
            for k in JuncDic:
                [ juncS, juncE ] = k.split( '_' )
                sortedJuncs.append( ( int( juncS ), int( juncE ), JuncDic[k] ) )
            sortedJuncs.sort\
                ( cmp = None, key = operator.itemgetter( 0 ), reverse = False )
            return sortedJuncs
        F_sorted_juncts = sort_junc_dict( self.F_Juncs )
        R_sorted_juncts = sort_junc_dict( self.R_Juncs )
        for J in F_sorted_juncts:
            print str( J[0] ), str( J[1] ), "%.4f" % round( J[2], 4 )

        for J in R_sorted_juncts:
            print str( J[0] ), str( J[1] ), "%.4f" % round( J[2], 4 )


